"""
可视化模型架构
"""
import torch
from config import Config
from model_definitions import create_model

def visualize_model_architecture():
    """可视化模型架构"""
    print("=== Transformer自编码器模型架构 ===")
    
    config = Config()
    model, criterion, optimizer = create_model(config)
    
    print(f"\n📊 **模型配置参数**")
    print(f"  输入维度: {config.CHANNEL_DIM} (64实部 + 64虚部)")
    print(f"  时间步长: {config.SEQ_LEN}")
    print(f"  批次大小: {config.BATCH_SIZE}")
    print(f"  端元数量: {config.K}")
    print(f"  稀疏度: {config.K_NONZERO}")
    print(f"  注意力头数: {config.NHEAD}")
    print(f"  Transformer层数: {config.NUM_LAYERS}")
    print(f"  Dropout率: {config.DROPOUT_RATE}")
    
    print(f"\n🏗️ **模型架构详解**")
    
    print(f"\n1️⃣ **输入层**")
    print(f"   输入形状: (batch_size, channel_dim, time_steps)")
    print(f"   实际形状: ({config.BATCH_SIZE}, {config.CHANNEL_DIM}, {config.SEQ_LEN})")
    print(f"   数据类型: 复数信号的实部虚部分离")
    
    print(f"\n2️⃣ **位置编码**")
    print(f"   类型: 可学习的位置编码")
    print(f"   形状: (1, batch_size, channel_dim)")
    print(f"   初始化: 正态分布 (mean=0, std=0.02)")
    print(f"   作用: 为时间序列添加位置信息")
    
    print(f"\n3️⃣ **多尺度注意力机制 (MultiScaleAttention)**")
    print(f"   🔍 自注意力: MultiheadAttention(d_model={config.CHANNEL_DIM}, nhead={config.NHEAD})")
    print(f"   🌊 多尺度卷积:")
    print(f"     - Conv1D(kernel=3): 捕获短期依赖")
    print(f"     - Conv1D(kernel=5): 捕获中期依赖") 
    print(f"     - Conv1D(kernel=7): 捕获长期依赖")
    print(f"     - Conv1D(kernel=9): 捕获超长期依赖")
    print(f"   ⚖️ 特征融合: 可学习权重加权融合5个特征")
    print(f"   🔄 残差连接: 输入 + 注意力输出")
    
    print(f"\n4️⃣ **Transformer编码器层 (PreNormTransformerEncoderLayer)**")
    print(f"   层数: {config.NUM_LAYERS + 1} 层")
    print(f"   结构: Pre-Norm架构 (先归一化再计算)")
    print(f"   子层1: 多尺度注意力 + 残差连接")
    print(f"   子层2: 前馈网络 + 残差连接")
    print(f"   前馈网络: {config.CHANNEL_DIM} → 1024 → {config.CHANNEL_DIM}")
    print(f"   激活函数: ReLU")
    print(f"   正则化: LayerNorm + Dropout({config.DROPOUT_RATE})")
    
    print(f"\n5️⃣ **输出层 (稀疏表示生成)**")
    print(f"   全连接网络:")
    print(f"     - Linear: {config.CHANNEL_DIM} → 512")
    print(f"     - ReLU + Dropout({config.DROPOUT_RATE})")
    print(f"     - Linear: 512 → {config.K}")
    print(f"     - ReLU + Dropout({config.DROPOUT_RATE})")
    print(f"   输出形状: (batch_size, k, time_steps)")
    
    print(f"\n6️⃣ **稀疏化处理**")
    print(f"   方法: TopK选择")
    print(f"   保留数量: 每个时间步保留前{config.K_NONZERO}个最大值")
    print(f"   归一化: L1归一化 (和为1)")
    print(f"   作用: 强制稀疏性约束")
    
    print(f"\n7️⃣ **重构层**")
    print(f"   方法: 矩阵乘法 Y = X @ D")
    print(f"   端元矩阵D: ({config.CHANNEL_DIM}, {config.K})")
    print(f"   稀疏表示X: (batch_size, time_steps, {config.K})")
    print(f"   重构输出: (batch_size, {config.CHANNEL_DIM}, {config.SEQ_LEN})")
    
    print(f"\n🎯 **损失函数 (CustomLoss)**")
    print(f"   重构损失: RMSE(Y_pred, Y_true)")
    print(f"   稀疏性损失: α × L1(X), α={config.ALPHA}")
    print(f"   平滑性损失: β × L1(X[:,:,1:] - X[:,:,:-1]), β={config.BETA}")
    print(f"   总损失: 重构损失 + 稀疏性损失 + 平滑性损失")
    
    print(f"\n⚙️ **优化器配置**")
    print(f"   优化器: Adam")
    print(f"   学习率: {config.LEARNING_RATE}")
    print(f"   权重衰减: {config.WEIGHT_DECAY}")
    
    # 计算模型参数
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"\n📈 **模型统计**")
    print(f"   总参数数量: {total_params:,}")
    print(f"   可训练参数: {trainable_params:,}")
    print(f"   模型大小: ~{total_params * 4 / 1024 / 1024:.2f} MB")
    
    # 测试前向传播
    print(f"\n🧪 **前向传播测试**")
    model.eval()
    with torch.no_grad():
        # 创建测试输入
        test_input = torch.randn(config.BATCH_SIZE, config.CHANNEL_DIM, config.SEQ_LEN).to(config.DEVICE)
        print(f"   输入形状: {test_input.shape}")
        
        # 前向传播
        Y_reconstructed, X = model(test_input)
        print(f"   重构输出形状: {Y_reconstructed.shape}")
        print(f"   稀疏表示形状: {X.shape}")
        
        # 检查稀疏性
        nonzero_count = torch.sum(X != 0, dim=1).float().mean()
        print(f"   平均非零元素数: {nonzero_count.item():.1f} (预期: {config.K_NONZERO})")
        
        # 计算损失
        loss = criterion(Y_reconstructed, test_input, X)
        print(f"   测试损失: {loss.item():.6f}")
    
    print(f"\n✅ **模型架构总结**")
    print(f"   这是一个基于Transformer的复数信号稀疏表示学习模型")
    print(f"   核心创新: 多尺度注意力机制 + 稀疏约束 + 时间序列建模")
    print(f"   应用场景: 复数信号的稀疏分解和重构")
    print(f"   优势: 能够捕获多时间尺度的依赖关系并强制稀疏性")
    
    return model, criterion, optimizer

if __name__ == "__main__":
    model, criterion, optimizer = visualize_model_architecture() 